import numpy as np
import mindspore as ms
import mindspore.nn as nn
import mindspore.dataset as ds

from util.callback import get_loss


class Callback(ms.Callback):
    def __init__(self, eval_data):
        self.eval_data = eval_data

    def epoch_end(self, run_context):
        callback_params = run_context.original_args()

        cur_epoch_num = callback_params.cur_epoch_num
        epoch_num = callback_params.epoch_num

        network = callback_params.network

        train_data = callback_params.train_dataset

        loss_t = get_loss(network, callback_params.loss_fn, train_data)
        loss_t /= callback_params.batch_num

        loss_v = get_loss(network, callback_params.loss_fn, self.eval_data)
        loss_v /= callback_params.batch_num

        print(f'[{cur_epoch_num}/{epoch_num}],train loss:{loss_t}, valid loss:{loss_v}')

        # print(f'[{cur_epoch_num}/{epoch_num}], W:{network[0].weight.asnumpy()}, b:{network[0].bias.asnumpy()}')


def get_data(num, d):
    true_w, true_b = np.ones((d, 1)) * 0.01, 0.05
    for _ in range(num):
        x = np.random.uniform(-10.0, 10.0, (1, d))
        noise = np.random.normal(0, 0.01)
        y = true_b + np.matmul(true_w, x) + noise
        yield np.array([x]).astype(np.float32), np.array([y]).astype(np.float32)


def create_dataset(num_data, batch_size, repeat_size):
    input_data = ds.GeneratorDataset(list(get_data(num_data, 100)), column_names=['data', 'label'])
    input_data = input_data.batch(batch_size)
    input_data = input_data.repeat(repeat_size)
    return input_data


def main():
    batch_size, lr, num_epochs, wd = 256, 0.0001, 10, 3

    dataset_train = create_dataset(20, 5, 1)
    dataset_test = create_dataset(100, 5, 1)

    net = nn.SequentialCell([nn.Dense(100, 1)])

    loss = nn.loss.MSELoss()
    opti = nn.SGD(net.trainable_params(), learning_rate=lr, weight_decay=wd)

    model = ms.Model(net, loss_fn=loss, optimizer=opti, metrics={'acc'})

    model.train(num_epochs, dataset_train, callbacks=[Callback(dataset_test)])


if __name__ == '__main__':
    main()

"""
output
[1/10],train loss:0.1743218, valid loss:1.5476905
[2/10],train loss:0.13951607, valid loss:1.4843119
[3/10],train loss:0.11297187, valid loss:1.4337035
[4/10],train loss:0.09191835, valid loss:1.3913558
[5/10],train loss:0.07553819, valid loss:1.3568121
[6/10],train loss:0.06261728, valid loss:1.3284906
[7/10],train loss:0.052319117, valid loss:1.3047441
[8/10],train loss:0.04395456, valid loss:1.2843126
[9/10],train loss:0.037215833, valid loss:1.2673154
[10/10],train loss:0.031726398, valid loss:1.2526271
"""
